# -*- coding: utf-8 -*-
"""
Created on Sat Sep 10 20:15:20 2022

@author: Shang'xiang

Stay Hungry Stay Foolish
"""

'''
此程序用于走时CT的正演
'''

import ttcrpy.rgrid as rg
import numpy as np
import matplotlib.pyplot as plt

# 创建网格
x = np.arange(0,31.0)
y = np.arange(0,21.0)

# 创建速度模型
v = 2000*np.ones((x.size-1,y.size-1))
v1 = 4000*np.ones((4, 4))

v[8:12,4:8] = v1
v[18:22,4:8] = v1

# 画模型图
fig, ax = plt.subplots()

cs = plt.pcolor(v,cmap='jet',edgecolor='w')
plt.colorbar()
plt.show()

# 给定发射点和接收点的坐标
nsrc = 30
nrcv = 30
srcs = np.ones(shape=(nsrc,2))
rcvs = 19*np.ones(shape=(nrcv,2))
for i in range(nsrc):
    xsrc = 0.5 + 1*i
    xrcv = 0.5 + 1*i
    srcs[i][0] = xsrc
    rcvs[i][0] = xrcv
    
# 离散网格
grid = rg.Grid2d(x, y, cell_slowness=True, method = 'SPM')

# 速度转换为慢度
slowness = 1./v

# tt_all = np.empty(0)
data_all = np.empty([900,5])
rays_all = list()
n = 0
for s in srcs:
    src = np.array([s])
    tt, rays = grid.raytrace(src, rcvs, slowness, return_rays=True)
    data_all[30*n:30*(n+1),0:2] = np.repeat(src,30,axis = 0)
    data_all[30*n:30*(n+1),2:4] = rcvs
    data_all[30*n:30*(n+1),4] = tt
    rays_all = rays_all + rays
    n = n + 1


# 绘制走时图
plt.figure(2)
plt.plot(data_all[:,4], 'r-o')
plt.show()

# 保存走时
np.savetxt('G03_FSM.txt', data_all)

plt.figure(3)
# 绘制射线路径图
for r in rays_all:
    plt.plot(r[:,1],r[:,0],'b-',linewidth=0.25)

ax = plt.gca()
miloc = plt.MultipleLocator(1)
ax.xaxis.set_minor_locator(miloc)
plt.grid(linestyle = ':', color = 'black',which = 'both')
plt.xlim(0,20)
plt.ylim(0,31)
plt.show()
